'''
实验6：图像复原
采用维纳滤波器对模糊图像进行复原，观察效果
对已知匀速直线运动模糊图像进行复原，观察效果
附加：任意运动模糊图像复原/图像去雾

利用OpenCV生成运动模糊图像和
再利用逆滤波和维纳滤波进行复原
'''

import numpy as np
from numpy import fft
import matplotlib.pyplot as plt
import math
import cv2

# 生成模糊矩阵（运动仿真向量）
def getMotionBlurKernel(degree, angle):
    Rotation = cv2.getRotationMatrix2D((degree / 2, degree / 2), angle, 1)
    kernelMotionBlur = np.zeros([degree,degree])
    for i in range(degree):
        kernelMotionBlur[int(degree/2)][i] = 1
    kernelMotionBlur = cv2.warpAffine(kernelMotionBlur, Rotation, (degree, degree))
    kernelMotionBlur = kernelMotionBlur / degree
    return kernelMotionBlur

# 利用模糊核产生模糊图像
def getimgBlur(img,kernelBlur):
    '''
    通过图像和模糊核在频域乘积进行模糊
    :param img:
    :param kernelBlur:模糊卷积核
    :return:
    '''
    imgHeight = img.shape[0]
    imgWdith = img.shape[1]
    kernelBlurHeight = kernelBlur.shape[0]
    kernelBlurWidth = kernelBlur.shape[1]
    newRow1 = np.zeros([(int)((imgHeight - kernelBlurHeight)/2), kernelBlurWidth])
    newRow2 = np.zeros([(imgHeight - kernelBlurHeight)-(int)((imgHeight - kernelBlurHeight)/2), kernelBlurWidth])
    newCol1 = np.zeros([imgHeight, (int)((imgWdith - kernelBlurWidth)/2)])
    newCol2 = np.zeros([imgHeight, (imgWdith - kernelBlurWidth)-(int)((imgWdith - kernelBlurWidth)/2)])
    kernelBlur = np.r_[newRow1,kernelBlur]
    kernelBlur = np.r_[kernelBlur,newRow2]
    kernelBlur = np.c_[newCol1,kernelBlur]
    kernelBlur = np.c_[kernelBlur,newCol2]

    img_fft = fft.fft2(img)
    kernelBlur_fft = fft.fft2(kernelBlur)
    imgBlur = np.abs(fft.fftshift(fft.ifft2(img_fft * kernelBlur_fft)))
    return imgBlur

# 逆滤波
def reverse(img,kernelBlur):
    '''
    逆滤波
    :param img:
    :param kernelBlur: 模糊卷积核
    :return:
    '''

    imgHeight = img.shape[0]
    imgWdith = img.shape[1]
    kernelBlurHeight = kernelBlur.shape[0]
    kernelBlurWidth = kernelBlur.shape[1]
    newRow1 = np.zeros([(int)((imgHeight - kernelBlurHeight) / 2), kernelBlurWidth])
    newRow2 = np.zeros([(imgHeight - kernelBlurHeight) - (int)((imgHeight - kernelBlurHeight) / 2), kernelBlurWidth])
    newCol1 = np.zeros([imgHeight, (int)((imgWdith - kernelBlurWidth) / 2)])
    newCol2 = np.zeros([imgHeight, (imgWdith - kernelBlurWidth) - (int)((imgWdith - kernelBlurWidth) / 2)])
    kernelBlur = np.r_[newRow1, kernelBlur]
    kernelBlur = np.r_[kernelBlur, newRow2]
    kernelBlur = np.c_[newCol1, kernelBlur]
    kernelBlur = np.c_[kernelBlur, newCol2]

    img_fft = fft.fft2(img)
    kernelBlur_fft = fft.fft2(kernelBlur)
    imgClear = np.abs(fft.fftshift(fft.ifft2(img_fft/kernelBlur_fft)))
    return imgClear

# 维纳滤波
def wiener(img,kernelBlur,k=0.01):
    '''
    维纳滤波
    :param img:
    :param kernelBlur:
    :param k: 维纳滤波系数k
    :return:
    '''

    imgHeight = img.shape[0]
    imgWdith = img.shape[1]
    kernelBlurHeight = kernelBlur.shape[0]
    kernelBlurWidth = kernelBlur.shape[1]
    newRow1 = np.zeros([(int)((imgHeight - kernelBlurHeight) / 2), kernelBlurWidth])
    newRow2 = np.zeros([(imgHeight - kernelBlurHeight) - (int)((imgHeight - kernelBlurHeight) / 2), kernelBlurWidth])
    newCol1 = np.zeros([imgHeight, (int)((imgWdith - kernelBlurWidth) / 2)])
    newCol2 = np.zeros([imgHeight, (imgWdith - kernelBlurWidth) - (int)((imgWdith - kernelBlurWidth) / 2)])
    kernelBlur = np.r_[newRow1, kernelBlur]
    kernelBlur = np.r_[kernelBlur, newRow2]
    kernelBlur = np.c_[newCol1, kernelBlur]
    kernelBlur = np.c_[kernelBlur, newCol2]

    img_fft = fft.fft2(img)
    kernelBlur_fft = fft.fft2(kernelBlur)
    imgClear = np.abs(fft.fftshift(fft.ifft2(np.abs(kernelBlur_fft)**2 * img_fft / kernelBlur_fft / (np.abs(kernelBlur_fft)**2 + k))))
    return imgClear


def blurTool(img, degree, angle):
    kernel = getMotionBlurKernel(degree, angle)
    imgBlur = np.abs(getimgBlur(img.copy(), kernel))
    imgBlur = np.clip(imgBlur, 0, 255).astype(np.uint8)
    return imgBlur.copy(), kernel.copy()

def wienerTools(img, kernelBlur, k=0.01):
    img_recover = wiener(img.copy(), kernelBlur.copy(), k=k)
    img_recover = np.clip(img_recover, 0, 255).astype(np.uint8)
    return img_recover.copy()


def reverseTools(img, kernelBlur):
    img_recover = reverse(img.copy(), kernelBlur.copy())
    img_recover = np.clip(img_recover, 0, 255).astype(np.uint8)
    return img_recover.copy()

if __name__ == '__main__':
    # 读图于与显示
    # 显示原图像
    image = cv2.imread('4g.jpg')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    plt.figure(1)
    plt.xlabel("Original Image")
    plt.gray()
    plt.xticks([])
    plt.yticks([])
    plt.imshow(image)

    # 参数设定
    degree=25           #模糊长度
    angle=30            #模糊角度
    k=0.01              #维纳滤波常数k
    percent=0.05        #高斯噪声百分比
    # 显示处理后的图像
    plt.figure(2)
    plt.gray()
    plt.suptitle("Blur Kernel:degree=%d,angle=%d"%(degree,angle))
    # 进行运动模糊处理
    kernel = getMotionBlurKernel(degree,angle)
    imgBlur = np.abs(getimgBlur(image, kernel))
    plt.subplot(231)
    plt.xlabel("Motion Blurred")
    plt.xticks([])
    plt.yticks([])
    plt.imshow(imgBlur)
    # 进行逆滤波
    result = reverse(imgBlur,kernel) #逆滤波
    plt.subplot(232)
    plt.xlabel("Reverse Deblurred")
    plt.xticks([])
    plt.yticks([])
    plt.imshow(result)
    # 进行维纳滤波
    result = wiener(imgBlur,kernel)
    plt.subplot(233)
    plt.xlabel("Wiener Deblurred,k=%.3f"%k)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(result)

    # 添加噪声
    imgBlurGuass = imgBlur + percent * imgBlur.std() * \
                   np.random.standard_normal(imgBlur.shape)
    plt.subplot(234)
    plt.xlabel("Motion & Noisy Blurred\npercent=%.3f"%percent)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(imgBlurGuass)
    # 对添加噪声的图像进行逆滤波
    plt.subplot(235)
    result = reverse(imgBlurGuass, kernel)
    plt.xlabel("Reverse Deblurred")
    plt.xticks([])
    plt.yticks([])
    plt.imshow(result)
    # 对添加噪声的图像进行维纳滤波
    result=wiener(imgBlurGuass,kernel)
    plt.subplot(236)
    plt.xlabel("Wiener Deblurred,k=%.3f"%k)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(result)

    # 显示两个模糊图像的傅里叶频谱，对数显示，中心化，归一化
    plt.figure(3)
    plt.gray()
    plt.suptitle("Blur Kernel:degree=%d,angle=%d"%(degree,angle))

    fouier_imgBlur = fft.fftshift(np.abs(np.log2(fft.fft2(imgBlur))) * 255.0 / np.abs(np.max(np.log2(fft.fft2(imgBlur)))))
    plt.subplot(121)
    plt.xlabel("imgBlur Fouier")
    plt.xticks([])
    plt.yticks([])
    plt.imshow(fouier_imgBlur)

    fouier_imgBlurGuass = fft.fftshift(np.abs(np.log2(fft.fft2(imgBlurGuass))) * 255.0 / np.abs(np.max(np.log2(fft.fft2(imgBlurGuass)))))
    plt.subplot(122)
    plt.xlabel("imgBlurGuass Fouier")
    plt.xticks([])
    plt.yticks([])
    plt.imshow(fouier_imgBlurGuass)

    plt.show()





''' 备注 '''
'''
1h
以下这串代码的意思是将kernelBlur矩阵扩展为和img矩阵同高宽的矩阵，并且位于新矩阵中心
imgHeight = img.shape[0]
imgWdith = img.shape[1]
kernelBlurHeight = kernelBlur.shape[0]
kernelBlurWidth = kernelBlur.shape[1]
newRow1 = np.zeros([(int)((imgHeight - kernelBlurHeight)/2), kernelBlurWidth])
newRow2 = np.zeros([(imgHeight - kernelBlurHeight)-(int)((imgHeight - kernelBlurHeight)/2), kernelBlurWidth])
newCol1 = np.zeros([imgHeight, (int)((imgWdith - kernelBlurWidth)/2)])
newCol2 = np.zeros([imgHeight, (imgWdith - kernelBlurWidth)-(int)((imgWdith - kernelBlurWidth)/2)])
kernelBlur = np.r_[newRow1,kernelBlur]
kernelBlur = np.r_[kernelBlur,newRow2]
kernelBlur = np.c_[newCol1,kernelBlur]
kernelBlur = np.c_[kernelBlur,newCol2]
'''
'''
OpenCV显示
fouier_imgBlur = fft.fftshift(np.abs(np.log2(fft.fft2(imgBlur))) * 255.0 / np.abs(np.max(np.log2(fft.fft2(imgBlur)))))
cv2.imshow("imgBlur Fouier",fouier_imgBlur.astype(np.uint8))
fouier_imgBlurGuass = fft.fftshift(np.abs(np.log2(fft.fft2(imgBlurGuass))) * 255.0 / np.abs(np.max(np.log2(fft.fft2(imgBlurGuass)))))
cv2.imshow("imgBlurGuass Fouier",fouier_imgBlurGuass.astype(np.uint8))
cv2.waitKey(0)
'''